package web

import (
	"bytes"
	"fmt"

	http "github.com/valyala/fasthttp"

	"xp/lib/errors"
	"xp/lib/mux"
	"xp/lib/mux/web/token"
	"xp/lib/utils"
)

const (
	// 为了简单的兼容我们直接使用angularjs的方式设置CSRF
	csrfTokenKey    = `XSRF-TOKEN`
	csrfReqTokenKey = `X-XSRF-TOKEN`
)

var (
	methodGet     = []byte(`GET`)
	htmlHeadTag   = []byte(`<head>`)
	bypassMethods = [][]byte{[]byte(`HEAD`), []byte(`OPTIONS`)}
)

func CSRF(next mux.Handler) mux.Handler {
	return func(ctx *mux.Ctx) {
		for _, method := range bypassMethods {
			if bytes.Equal(method, ctx.Method()) {
				next(ctx)
				return
			}
		}

		// 只在GET方法的时候生成Token
		if bytes.Equal(ctx.Method(), methodGet) && !IsAjax(ctx) {
			next(ctx)

			// 非HTML响应我们不修改CSRF-TOKEN
			if !bytes.Equal(ctx.Response.Header.ContentType(), contentTypeHtml) {
				return
			}

			// 每次使用新的CRSF TOKEN是为了防止TLS + GZIP的情况可能攻击的情况
			tk, err := token.New(ctx, token.CurrentUser(ctx))
			if err != nil {
				ctx.Error(err)
				return
			}

			// 替换html中的head标签，添加CSRF-TOKEN
			html := ctx.Response.Body()
			csrfTokenMeta := fmt.Sprintf(`<meta name="%s" content="%s">`, csrfTokenKey, tk)
			html = bytes.Replace(html, htmlHeadTag, append(htmlHeadTag, csrfTokenMeta...), 1)
			ctx.Response.SetBody(html)

			// session cookie
			cookie := http.AcquireCookie()
			cookie.SetKey(csrfTokenKey)
			cookie.SetValue(tk)
			cookie.SetPath("/")
			host := ctx.Host()
			if index := bytes.Index(host, []byte(":")); index > 0 {
				host = host[:index]
			}
			cookie.SetDomainBytes(host)
			ctx.Response.Header.SetCookie(cookie)
			// disable cache
			ctx.Response.Header.Set("Cache-Control", "no-cache, no-store, must-revalidate")
			http.ReleaseCookie(cookie)
			return
		}

		tk := utils.Bytes2Str(ctx.GetHeader(csrfReqTokenKey))
		pass, err := token.Validate(ctx, tk, token.CurrentUser(ctx))
		if err != nil {
			ctx.Error(err)
			return
		}
		if !pass {
			ctx.Error(errors.CodeForbidden)
			return
		}

		cookieToken := utils.Bytes2Str(ctx.Request.Header.Cookie(csrfTokenKey))
		if tk != cookieToken {
			// 处理多个服务同时设置同样的cookie的情况
			pass = false
			csrfTokenKeyBytes := utils.Str2Bytes(csrfTokenKey)
			ctx.Request.Header.VisitAllCookie(func(key, value []byte) {
				if !bytes.Equal(key, csrfTokenKeyBytes) {
					return
				}
				if tk == utils.Bytes2Str(value) {
					pass = true
				}
			})
		}
		if !pass {
			ctx.Error(errors.CodeForbidden)
			return
		}

		next(ctx)
	}
}
